import torch
import torch.onnx
import sys
sys.path.append("/home/autox/anaconda3/python_learn/tutorial/class_1")
import netualbuild

model = torch.load("NeuralNetwork.pth", weights_only=False)
model.eval()
print(model)

batch_size=1
dummy_input=torch.randn(batch_size,1,28,28,requires_grad=True)
dummy_input = dummy_input.to("cuda:0")
print(dummy_input.shape)
out_put=model(dummy_input)
torch.onnx.export(model,dummy_input,"NeuralNetwork.onnx",export_params=True)